//=================================================================================================
// Copyright (c) 2011, Stefan Kohlbrecher, TU Darmstadt
// All rights reserved.

// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above copyright
//       notice, this list of conditions and the following disclaimer in the
//       documentation and/or other materials provided with the distribution.
//     * Neither the name of the Simulation, Systems Optimization and Robotics
//       group, TU Darmstadt nor the names of its contributors may be used to
//       endorse or promote products derived from this software without
//       specific prior written permission.

// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
// ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER BE LIABLE FOR ANY
// DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
// ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//=================================================================================================

#ifndef GOOGLE_MAPPING_SCAN_MATCHER_H_
#define GOOGLE_MAPPING_SCAN_MATCHER_H_

#include <cmath>
#include <Eigen/Geometry>

#include "Map.h"

namespace hectorslam {

class MotionModel {
 public:
  virtual void getPredictedPose(const ros::Time& time, tf::Pose* result) const = 0;
  virtual bool poseDifferenceTooLarge(const ros::Time& time, const tf::Pose& pose) const = 0;
  virtual void addPoseObservation(const ros::Time& time, const tf::Pose& observation) = 0;
};

class DumbMotionModel : public MotionModel {
 public:
  virtual void getPredictedPose(const ros::Time& time, tf::Pose* result) const { *result = pose_; }
  virtual bool poseDifferenceTooLarge(const ros::Time& time, const tf::Pose& pose) const { return false; }
  virtual void addPoseObservation(const ros::Time& time, const tf::Pose& observation) {
    pose_ = observation;
  }
 private:
  tf::Pose pose_;
};

class HumanMotionModel : public MotionModel {
 public:
  HumanMotionModel(double linear_speed_threshold, double angular_speed_threshold)
     : linear_speed_threshold_(linear_speed_threshold),
       angular_speed_threshold_(angular_speed_threshold),
       pose_(tf::Pose::getIdentity()){}

  virtual void getPredictedPose(
      const ros::Time& time, tf::Pose* result) const {
    *result = pose_;
  }

  virtual bool poseDifferenceTooLarge(const ros::Time& time, const tf::Pose& pose) const {
    double distance = pose_.getOrigin().distance(pose.getOrigin());
    double angle = abs(pose_.getRotation().angle(pose.getRotation()));
    double max_distance = linear_speed_threshold_ * (time - time_).toSec();
    double max_angle = angular_speed_threshold_ * (time - time_).toSec();
    ROS_INFO("dt: %f, distance: %f, max_distance: %f, angle: %f, max_angle: %f", (time-time_).toSec(), distance, max_distance, angle, max_angle);
    return distance > max_distance || angle > max_angle;
  }

  virtual void addPoseObservation(const ros::Time& time, const tf::Pose& observation) {
    time_ = time;
    pose_ = observation;
  }

 private:
  ros::Time time_;
  tf::Pose pose_;
  double linear_speed_threshold_;
  double angular_speed_threshold_;
};

class ScanMatcher {
 public:
  EIGEN_MAKE_ALIGNED_OPERATOR_NEW

  bool match(
      const google_mapping::Map& map,
      const pcl::PointCloud<pcl::PointXYZ>& point_cloud,
      const MotionModel& motion_model,
      const ros::Time& time,
      tf::Pose* result) {
    tf::Pose initial_pose_estimate;
    motion_model.getPredictedPose(time, &initial_pose_estimate);
    btScalar unused_roll, unused_pitch, yaw;
    tf::Quaternion rotation = initial_pose_estimate.getRotation();
    tf::Matrix3x3(rotation).getRPY(unused_roll, unused_pitch, yaw);
    Eigen::Vector3f pose_estimate(initial_pose_estimate.getOrigin().x(),
                                  initial_pose_estimate.getOrigin().y(),
                                  yaw);
    for (int level = 2; level >= 0; --level) {
      double start_fitness = computePointCloudFitness(
          map, point_cloud, pose_estimate, level);
      int maxIterations = 2;
      bool stopped = false;
      Eigen::Vector3f optimized_pose_estimate = pose_estimate;
      for (int j = 0; j < maxIterations && !stopped; ++j) {
        stopped = estimateTransformationLogLh(
            map, point_cloud, optimized_pose_estimate, level, &optimized_pose_estimate);
      }
      if (motion_model.poseDifferenceTooLarge(time, createPose(optimized_pose_estimate))) {
        std::cout<<"Moved too far: " << level << "\n";
        continue;
      }
      double final_fitness = computePointCloudFitness(
          map, point_cloud, optimized_pose_estimate, level);
      if (final_fitness - start_fitness > 1e-4) {
        return false;
      }
      if (final_fitness < start_fitness) {
        pose_estimate = optimized_pose_estimate;
      }
    }
    *result = createPose(pose_estimate);
    return true;
  }

 private:
  tf::Pose createPose(Eigen::Vector3f vector) {
    vector.z() = normalize_angle(vector.z());
    double w = cos(vector.z() * 0.5f);
    double z = sin(vector.z() * 0.5f);
    return tf::Pose(
        tf::Quaternion(0., 0., z, w),
        tf::Vector3(vector.x(), vector.y(), 0.));
  }

  inline float normalize_angle(float angle) {
    float a = fmod(fmod(angle, 2.f * M_PI) + 2.f * M_PI, 2.f * M_PI);
    if (a > M_PI) {
      a -= 2.f * M_PI;
    }
    return a;
  }

  bool estimateTransformationLogLh(
      const google_mapping::Map& map,
      const pcl::PointCloud<pcl::PointXYZ>& point_cloud,
      const Eigen::Vector3f& initial_pose_estimate,
      int level,
      Eigen::Vector3f* result) {
    *result = initial_pose_estimate;
    computeHessianAndGradient(map, point_cloud, initial_pose_estimate, level);
    float unusedDeterminant;
    bool invertible;
    hessian_.computeInverseAndDetWithCheck(
        hessianInverse_, unusedDeterminant, invertible);
    if (!invertible) {
      std::cerr << "Hessian not invertible." << std::endl;
      return true;
    }
    *result += hessianInverse_ * gradient_;
    return gradient_.norm() < 1e-4;
  }

  double computePointCloudFitness(
      const google_mapping::Map& map,
      const pcl::PointCloud<pcl::PointXYZ>& point_cloud,
      const Eigen::Vector3f& pose_estimate,
      int level) {
    Eigen::Affine2f transform(
        Eigen::Translation2f(pose_estimate[0], pose_estimate[1])
        * Eigen::Rotation2Df(pose_estimate[2]));
    double global_fitness = 0.;
    for (size_t i = 0; i < point_cloud.size(); ++i) {
      Eigen::Vector2f transformed_point =
          transform * Eigen::Vector2f(point_cloud[i].x, point_cloud[i].y);
      Eigen::Vector3f values = map.get2DInterpolatedProbabilityWithDerivatives(
          transformed_point.x(), transformed_point.y(), point_cloud[i].z, level);
      double local_fitness = 1. - values[0];
      global_fitness += local_fitness * local_fitness;
    }
    return global_fitness;
  }
 
  float computeHessianAndGradient(
      const google_mapping::Map& map,
      const pcl::PointCloud<pcl::PointXYZ>& point_cloud,
      const Eigen::Vector3f& start_estimate,
      int level) {
    Eigen::Affine2f transform(
        Eigen::Translation2f(start_estimate[0], start_estimate[1])
        * Eigen::Rotation2Df(start_estimate[2]));

    double sin_rotation = sin(start_estimate[2]);
    double cos_rotation = cos(start_estimate[2]);

    double global_fitness = 0.;
    hessian_ = Eigen::Matrix3f::Zero();
    gradient_ = Eigen::Vector3f::Zero();

    for (size_t i = 0; i < point_cloud.size(); ++i) {
      Eigen::Vector2f point =
          Eigen::Vector2f(point_cloud[i].x, point_cloud[i].y);
      Eigen::Vector2f transformed_point = transform * point;
      Eigen::Vector3f values = map.get2DInterpolatedProbabilityWithDerivatives(
          transformed_point.x(), transformed_point.y(), point_cloud[i].z, level);

      double local_fitness = 1. - values[0];
      global_fitness += local_fitness * local_fitness;

      gradient_[0] += values[1] * local_fitness;
      gradient_[1] += values[2] * local_fitness;

      double rotational_derivative = 
          (-sin_rotation * point.x() - cos_rotation * point.y()) * values[1] + 
          (cos_rotation * point.x() - sin_rotation * point.y()) * values[2];

      gradient_[2] += rotational_derivative * local_fitness;

      hessian_(0, 0) += std::pow(values[1], 2);
      hessian_(1, 1) += std::pow(values[2], 2);
      hessian_(2, 2) += std::pow(rotational_derivative, 2);

      hessian_(0, 1) += values[1] * values[2];
      hessian_(0, 2) += values[1] * rotational_derivative;
      hessian_(1, 2) += values[2] * rotational_derivative;
    }

    hessian_(1, 0) = hessian_(0, 1);
    hessian_(2, 0) = hessian_(0, 2);
    hessian_(2, 1) = hessian_(1, 2);

    return global_fitness;
  }

  Eigen::Vector3f gradient_;
  Eigen::Matrix3f hessian_;
  Eigen::Matrix3f hessianInverse_;
};

}  // namespace hectorslam

#endif  // GOOGLE_MAPPING_SCAN_MATCHER_H_
